# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import math
from collections import deque

import numpy as np


class MovingWindow(object):
  def __init__(self, timewindow):
    self._timewindow = timewindow
    self._dq = deque()
    self._last_popped = None

  def get_last_popped(self):
    return self._last_popped

  def iteritems(self):
    for ts, obj in self._dq:
      yield ts, obj

  def update(self, ts, data=None):
    raise NotImplementedError()

  def get_latest_data(self):
    raise NotImplementedError()

  def get_first_data(self):
    raise NotImplementedError()

  def ready(self):
    return self._dq and self._dq[-1][0] - self._dq[0][0] > 0.3 * self._timewindow

  def size(self):
    return len(self._dq)

  def get_latest_ts(self):
    if self._dq:
      return self._dq[-1][0]
    return 0


class MovingSum(MovingWindow):
  def __init__(self, timewindow):
    MovingWindow.__init__(self, timewindow)
    self._multiplier = 1e6
    self._imultiplier = 1e-6
    self._sum = 0

  def update(self, ts, data=None):
    while self._dq and ts - self._dq[0][0] > self._timewindow:
      self._sum -= self._dq[0][1]
      self._last_popped = self._dq[0]
      self._dq.popleft()
    if data is not None:
      data = int(data * self._multiplier + 0.4999)
      self._dq.append((ts, data))
      self._sum += data

  def get_latest_data(self):
    if self._dq:
      return self._dq[-1][1] * self._imultiplier
    else:
      return np.nan

  def get_statistics(self):
    return self._sum * self._imultiplier


class MovingAverage(MovingWindow):
  def __init__(self, timewindow):
    MovingWindow.__init__(self, timewindow)
    self._multiplier = 1e6
    self._imultiplier = 1e-6
    self._sum = 0
    self._sum_sqr = 0

  def update(self, ts, data=None):
    while self._dq and ts - self._dq[0][0] > self._timewindow:
      self._sum -= self._dq[0][1]
      self._sum_sqr -= self._dq[0][2]
      self._last_popped = self._dq[0]
      self._dq.popleft()
    if data is not None:
      data = int(data * self._multiplier + 0.4999)
      data2 = data * data
      self._dq.append((ts, data, data2))
      self._sum += data
      self._sum_sqr += data2

  def get_latest_data(self):
    if self._dq:
      return self._dq[-1][1] * self._imultiplier
    else:
      return np.nan

  def get_statistics(self):
    n = len(self._dq)
    mean = 1.0 * self._sum / n
    std = 1.0 * (self._sum_sqr - n * mean * mean) / n
    return mean * self._imultiplier, math.sqrt(std) * self._imultiplier


class TimeWeightedMovingAverage(MovingWindow):
  def __init__(self, timewindow):
    MovingWindow.__init__(self, timewindow)
    self._multiplier = 1e6
    self._imultiplier = 1e-6
    self._sum = 0

  def update(self, ts, data=None):
    ts = int(ts)
    while self._dq and ts - self._dq[0][0] > self._timewindow:
      if len(self._dq) > 1:
        self._sum -= (self._dq[1][0] - self._dq[0][0]) * self._dq[0][1]
      else:
        assert self._sum == 0, self._sum
      self._last_popped = self._dq[0]
      self._dq.popleft()
    if data is not None:
      data = int(data * self._multiplier + 0.4999)
      if self._dq:
        self._sum += (ts - self._dq[-1][0]) * self._dq[-1][1]
      self._dq.append((ts, data))

  def get_latest_data(self):
    if self._dq:
      return self._dq[-1][1] * self._imultiplier
    else:
      return np.nan

  def get_statistics(self, ts=None):
    if self._dq:
      summ = self._sum
      tss = self._dq[-1][0] - self._dq[0][0]
      if ts and ts != self._dq[-1][0]:
        dts = ts - self._dq[-1][0]
        summ += self._dq[-1][1] * dts
        tss += dts
      if tss > 0:
        return summ / tss * self._imultiplier
    return np.nan


class VolumeWeightedMovingAverage(MovingWindow):
  def __init__(self, timewindow):
    MovingWindow.__init__(self, timewindow)
    self._multiplier = 1e5
    self._imultiplier = 1e-5
    self._sum_pq = 0
    self._sum_q = 0

  def update(self, ts, pq=None):
    ts = int(ts)
    while self._dq and ts - self._dq[0][0] > self._timewindow:
      self._sum_pq -= self._dq[0][1]
      self._sum_q -= self._dq[0][2]
      self._last_popped = self._dq[0]
      self._dq.popleft()
    if pq is not None:
      p, q = pq
      pxq = int(p * q * self._multiplier + 0.4999)
      q = int(q * self._multiplier + 0.4999)
      self._sum_pq += pxq
      self._sum_q += q
      self._dq.append((ts, pxq, q))

  def get_statistics(self, ts=None):
    if ts:
      self.update(ts)
    if self._sum_q > 0:
      return self._sum_pq / self._sum_q
    return np.nan


class MovingObject(MovingWindow):
  def __init__(self, timewindow):
    MovingWindow.__init__(self, timewindow)

  def update(self, ts, data=None):
    while self._dq and ts - self._dq[0][0] > self._timewindow:
      self._last_popped = self._dq[0]
      self._dq.popleft()
    if data is not None:
      self._dq.append((ts, data))

  def get_latest_data(self):
    if self._dq:
      return self._dq[-1][1]
    else:
      return None

  def get_first_data(self):
    if self._dq:
      return self._dq[0][1]
    else:
      return None
